/*
 * Copyright (C) 2007 Eskil Bylund
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Threading;

using DCSharp.Backend.Connections;
using DCSharp.Backend.Objects;
using DCSharp.Hashing;
using DCSharp.Logging;
using DCSharp.Security.Cryptography;

namespace DCSharp.Backend.Managers
{
	/// <summary>
	/// The errors that can happen when downloading from a <see cref="User"/>.
	/// </summary>
	public enum FileDownloadError
	{
		None,
		FileNotAvailable,
		ClientToOld,
		InvalidHashTree,
		ValidationError
	}

	/// <summary>
	/// Handles the downloading of a single file.
	/// </summary>
	/// <remarks>
	/// The steps involved to successfully download a file are:
	///
	/// 1. Get the hash tree. The first time it's downloaded from a user, after
	///    that it's obtained from the cache. The hash tree is used to verify
	///    the temp file and all the downloaded data.
	///
	/// 2. Verify the temp file. A BitArray is used to keep track of the finished
	///    and unfinished segments, and is updated when the temp file is verified
	///    and when data is received.
	///
	/// 3. Download the file data. The ISegmentStrategy determines what segments
	///    to use when downloading.
	///
	/// 4. Move the temp file to the correct location. Done!
	/// </remarks>
	public class FileDownloader : Downloader, IEnumerable<SourceInfo>
	{
		private static Logger log = LogManager.GetLogger("Downloader");

		public event EventHandler<SourceEventArgs> SourceAdded;
		public event EventHandler<SourceEventArgs> SourceRemoved;
		private delegate void HashHandler(object state);

		private static ActionQueue<Downloader> hashingQueue = new ActionQueue<Downloader>("Download Verification");
		private HashStore hashStore;

		private RemoteFileInfo remoteFile;
		private string savePath;
		private long downloadedBytes;
		private Dictionary<User, SourceInfo> sources;
		private List<SourceInfo> sourcePriority; // The priority is determined by the position in this list

		private HashTree hashTree; // The hash tree used for verification
		private ISegmentStrategy segmentStrategy; // Responsible for creating segments
		private SegmentedFileInfo segmentInfo; // Holds information needed when creating segments. Created after the temp file's been verified.

		// TODO: Use a list of users to keep track of the user requests instead
		// of the int pendingRequestCount?

		private Dictionary<UserConnection, WorkerState> connections; // The active connections
		private int pendingRequestCount; // Represents the number of pending user requests
		private User pendingHashTreeUser; // The user the hash tree is set to be downloaded from
		private bool handleUserIdle; // Used by ForceStart to get the user requests in the correct order

		private FileStream outputStream;
		private object streamLock; // The lock for the output stream

		public FileDownloader(RemoteFileInfo remoteFile, string savePath,
			HashStore hashStore) : this(remoteFile, savePath, hashStore,
			new MultiSegmentStrategy())
		{
		}

		public FileDownloader(RemoteFileInfo remoteFile, string savePath,
			HashStore hashStore, ISegmentStrategy strategy)
		{
			if (remoteFile == null)
			{
				throw new ArgumentNullException("remoteFile");
			}
			if (savePath == null)
			{
				throw new ArgumentNullException("savePath");
			}
			if (hashStore == null)
			{
				throw new ArgumentNullException("hashStore");
			}
			if (strategy == null)
			{
				throw new ArgumentNullException("strategy");
			}
			this.remoteFile = remoteFile;
			this.savePath = savePath;
			this.hashStore = hashStore;
			this.segmentStrategy = strategy;

			sources = new Dictionary<User, SourceInfo>();
			sourcePriority = new List<SourceInfo>();

			connections = new Dictionary<UserConnection, WorkerState>();
			handleUserIdle = true;
			streamLock = new Object();

			hashTree = hashStore.GetHashTree(remoteFile.TTH);
		}

		#region Classes

		/// <summary>
		/// Contains the objects used when downloading a segment.
		/// </summary>
		private class WorkerState
		{
			public UserConnection Connection;
			public SourceInfo Source;
			public Segment Segment;
			public Stream Stream; // The output stream

			public WorkerState(UserConnection connection, SourceInfo source,
				Segment segment, Stream stream)
			{
				this.Connection = connection;
				this.Source = source;
				this.Segment = segment;
				this.Stream = stream;
			}

			public override string ToString()
			{
				return Source.User.ToString();
			}
		}

		#endregion

		#region Properties

		public RemoteFileInfo RemoteFile
		{
			get { return remoteFile; }
		}

		public string SavePath
		{
			get { return savePath; }
		}

		public override long FileSize
		{
			get { return remoteFile.Size; }
		}

		public ISegmentStrategy SegmentStrategy
		{
			get { return segmentStrategy; }
			set { segmentStrategy = value; }
		}

		public override object UserLock
		{
			get { return sources; }
		}

		internal override IEnumerable<User> InterestingUsers
		{
			get { return sources.Keys; }
		}

		#endregion

		#region Methods

		#region Source

		public void AddSource(User user)
		{
			if (user == null)
			{
				throw new ArgumentNullException("user");
			}

			SourceInfo source = new SourceInfo(user);
			lock (sources)
			{
				if (sources.ContainsKey(user))
				{
					return;
				}
				sources.Add(user, source);
				sourcePriority.Add(source);
			}
			OnSourceAdded(source);
			OnUserInteresting(user);
		}

		public void RemoveSource(User user)
		{
			if (user == null)
			{
				throw new ArgumentNullException("user");
			}

			SourceInfo source;
			lock (StateLock)
			{
				lock (sources)
				{
					if (!sources.TryGetValue(user, out source))
					{
						return;
					}
					sources.Remove(user);
					sourcePriority.Remove(source);
				}

				// Disconnect if downloading from the source
				foreach (KeyValuePair<UserConnection, WorkerState> pair in connections)
				{
					if (pair.Value.Source == source)
					{
						pair.Value.Connection.Disconnect();
						break;
					}
				}
			}
			OnSourceRemoved(source);
			OnUserUninteresting(source.User);
		}

		public bool ContainsSource(User user)
		{
			return sources.ContainsKey(user);
		}

		public IEnumerator<SourceInfo> GetEnumerator()
		{
			return sources.Values.GetEnumerator();
		}

		IEnumerator IEnumerable.GetEnumerator()
		{
			return GetEnumerator();
		}

		protected virtual void OnSourceAdded(SourceInfo source)
		{
			if (SourceAdded != null)
			{
				SourceAdded(this, new SourceEventArgs(source));
			}
		}

		protected virtual void OnSourceRemoved(SourceInfo source)
		{
			if (SourceRemoved != null)
			{
				SourceRemoved(this, new SourceEventArgs(source));
			}
		}

		#endregion

		public override void Start()
		{
			lock (StateLock)
			{
				if (State == DownloadState.Stopped)
				{
					log.Info("Starting", null, this);
					if (hashTree != null && segmentInfo == null)
					{
						CheckTempFile();
					}
					else
					{
						State = DownloadState.Idle;
						RequestConnections();
					}
				}
			}
		}

		public override void Stop()
		{
			lock (StateLock)
			{
				hashingQueue.Abort(this);
				base.Stop();

				foreach (UserConnection connection in
					new List<UserConnection>(connections.Keys))
				{
					connection.Disconnect();
				}
				lock (streamLock)
				{
					CloseOutputStream();
				}
			}
		}

		public void RemoveTempFile()
		{
			lock (StateLock)
			{
				Stop();

				string tempFile = GetTempFileName();
				if (File.Exists(tempFile))
				{
					File.Delete(tempFile);
				}
			}
		}

		public void ForceStart(SourceInfo source)
		{
			lock (sources)
			{
				if (!sources.ContainsValue(source))
				{
					return;
				}
				sourcePriority.Remove(source);
				sourcePriority.Insert(0, source);
			}
			lock (StateLock)
			{
				// FIXME: The forced start will not work if the DownloadManager
				// has a downloader with a higher priority.
				int availableSegments = segmentInfo != null ?
					segmentStrategy.GetAvailableSegments(segmentInfo) : -1;
				if (availableSegments == 0 && connections.Count > 0)
				{
					// Disconnect the slowest connection
					UserConnection slowest = null;
					foreach (UserConnection c in connections.Keys)
					{
						if (slowest == null ||
							c.Monitor.DownloadSpeed < slowest.Monitor.DownloadSpeed)
						{
							slowest = c;
						}
					}
					if (slowest != null)
					{
						// Ignore any calls to HandleUserIdle. RequestConnections
						// will be called by OnTransferAborted/Reset.
						handleUserIdle = false;

						slowest.Disconnect();

						handleUserIdle = true;
					}
				}
				else
				{
					RequestConnections(source.User);
				}
			}
		}

		public override string ToString()
		{
			return remoteFile.ToString();
		}

		#region Request

		protected override void RequestConnections()
		{
			RequestConnections(null);
		}

		internal override void HandleIdleUser(User user)
		{
			if (handleUserIdle)
			{
				RequestConnections(user);
			}
		}

		/// <summary>
		/// Request a connection for each available segment, or request
		/// a single connection if the hash tree needs to be downloaded.
		/// </summary>
		private void RequestConnections(User targetUser)
		{
			lock (StateLock)
			{
				if (State != DownloadState.Idle && State != DownloadState.Active)
				{
					return;
				}
				if (hashTree == null && (pendingHashTreeUser != null ||
					connections.Count > 0))
				{
					// Pending/active hash tree download found.
					return;
				}
				else if (hashTree != null && segmentInfo == null)
				{
					// The temp file needs to be verified.
					return;
				}
				lock (sources)
				{
					foreach (SourceInfo source in sourcePriority)
					{
						if (!source.IsValid ||
							(targetUser != null && source.User != targetUser))
						{
							continue;
						}
						if (segmentInfo != null)
						{
							int availableSegments = segmentStrategy.GetAvailableSegments(segmentInfo);
							if (availableSegments == 0 || pendingRequestCount >= availableSegments)
							{
								return;
							}
						}
						if (OnConnectionRequest(source.User, RequestCallback))
						{
							log.Info("Requested a connection to " + source.User, null, this);
							Interlocked.Increment(ref pendingRequestCount);

							// Done?
							if (hashTree == null)
							{
								pendingHashTreeUser = source.User;
								return;
							}
						}
					}
				}
			}
		}

		private void RequestCallback(User user, bool success)
		{
			Interlocked.Decrement(ref pendingRequestCount);
			pendingHashTreeUser = null;

			if (!success)
			{
				log.Error("Connection request to " + user + " failed", null, this);
				RequestConnections();
			}
		}

		#endregion

		#region TempFile

		// Hold before calling: StateLock
		private void CheckTempFile()
		{
			Debug.Assert(hashTree != null);

			State = DownloadState.Loading;
			try
			{
				// If the temp file doesn't exist, run the check directly so
				// that we don't have to wait for the other downloaders to finish.
				if (File.Exists(GetTempFileName()))
				{
					hashingQueue.Queue(this, PerformHashCheck);
				}
				else
				{
					PerformHashCheck(null);
				}
			}
			catch (Exception e)
			{
				log.Error("Error accessing the temp file.", e, this);
				StopWithException(e);
			}
		}

		// A SegmentManager will be created if the verification is successful.
		private void PerformHashCheck(object state)
		{
			Debug.Assert(hashTree != null);

			log.Info("Checking temp file", null, this);
			try
			{
				// Open the temp file
				lock (streamLock)
				{
					CloseOutputStream();
					OpenOutputStream(FileAccess.Read);

					if (outputStream.Length > remoteFile.Size)
					{
						outputStream.SetLength(remoteFile.Size);
					}
				}

				downloadedBytes = 0;

				// Verify the data
				HashTree tempTree = new HashTree(Tiger.Create(), hashTree.BlockSize);
				BitArray completed = new BitArray(hashTree.LeafCount);
				int leafPos = 0;
				do
				{
					lock (streamLock)
					{
						if (outputStream != null)
						{
							tempTree.Update(outputStream, hashTree.BlockSize);
							if (outputStream.Position == remoteFile.Size)
							{
								tempTree.CalculateRoot();
							}
						}
					}

					// Since the tree is updated with data for a single block,
					// we know that only one or zero blocks are created.
					if (leafPos < hashTree.LeafCount &&
						leafPos < tempTree.LeafCount)
					{
						completed[leafPos] = HashTree.LeavesAreEqual(
							tempTree[leafPos], hashTree[leafPos]);

						if (completed[leafPos])
						{
							downloadedBytes += hashTree.BlockSize;
							Progress = (double)downloadedBytes / remoteFile.Size;
						}
						leafPos++;
					}
					lock (streamLock)
					{
						if (outputStream == null ||
							outputStream.Position == outputStream.Length ||
							hashingQueue.AbortEvent.WaitOne(0, false))
						{
							break;
						}
					}
				}
				while (true);

				// Finish the verification
				if (downloadedBytes > remoteFile.Size)
				{
					downloadedBytes = remoteFile.Size;
					Progress = 1;
				}
				if (downloadedBytes == remoteFile.Size)
				{
					lock (StateLock)
					{
						HandleDownloadComplete();
					}
				}
				else if (!hashingQueue.AbortEvent.WaitOne(0, false))
				{
					int leavesDone = (int)((downloadedBytes + hashTree.BlockSize - 1) /
						hashTree.BlockSize);
					log.Info(leavesDone + " of " + hashTree.LeafCount + " leaves done",
						null, this);

					segmentInfo = new SegmentedFileInfo(remoteFile, hashTree,
						completed, new BitArray(completed.Length));
				}
			}
			catch (Exception e)
			{
				log.Error("Checking the temp file failed.", e, this);
				StopWithException(e);
			}
			finally
			{
				lock (streamLock)
				{
					CloseOutputStream();
				}
				lock (StateLock)
				{
					if (State == DownloadState.Loading)
					{
						State = DownloadState.Idle;
					}
					RequestConnections();
				}
			}
		}

		#endregion

		#region Download

		internal override bool Use(UserConnection connection)
		{
			lock (StateLock)
			{
				if (State != DownloadState.Idle && State != DownloadState.Active)
				{
					return false;
				}
				if ((hashTree == null && connections.Count > 0) ||
					(hashTree != null && segmentInfo == null) ||
					connections.ContainsKey(connection))
				{
					return false;
				}
				lock (sources)
				{
					SourceInfo source;
					if (sources.TryGetValue(connection.RemoteUser, out source) &&
						source.IsValid)
					{
						if (hashTree == null)
						{
							pendingHashTreeUser = null;
							return DownloadHashTree(connection, source);
						}
						else if (segmentInfo != null)
						{
							return DownloadFile(connection, source);
						}
					}
					return false;
				}
			}
		}

		private void UpdateDownloadPosition(int received)
		{
			downloadedBytes += received;
			Progress = (double)downloadedBytes / remoteFile.Size;
		}

		// NOTE: Don't call reset with the same state multiple times, since the
		// unvalidated bytes will be subtracted from the download position each time.
		// Hold before calling: StateLock
		private void Reset(WorkerState state)
		{
			log.Debug("Resetting " + state, null, this);

			Segment segment = state.Segment;
			if (segment != null)
			{
				UpdateBitArray(segmentInfo.Active, segment, false);

				// Revert any bytes not validated
				ValidatingOutputStream stream = state.Stream as ValidatingOutputStream;
				if (stream != null && segment.Position < segment.End)
				{
					long validated = hashTree.BlockSize * stream.ValidatedLeaves;
					long notValidated = (segment.Position - segment.Start) - validated;

					lock (streamLock)
					{
						UpdateDownloadPosition((int)-notValidated);
					}
				}
			}

			state.Connection.TransferAborted -= OnTransferAborted;
			state.Source.Monitor = null;
			state.Source.Active = false;

			connections.Remove(state.Connection);
			if (connections.Count == 0 && State == DownloadState.Active)
			{
				log.Debug("No active segments. Changing state to Idle.", null, this);
				State = DownloadState.Idle;
			}
			// FIXME: Really request connections from Reset?
			RequestConnections();
		}

		private void OnTransferAborted(object obj, TransferErrorEventArgs args)
		{
			WorkerState state;
			if (!connections.TryGetValue((UserConnection)obj, out state))
			{
				return;
			}

			log.Error("Transfer from " + state.Source.User + " aborted: " + args.Error + " - " + args.Message,
				null, this);

			if (args.Error == TransferError.NotAvailable)
			{
				SetValid(state.Source, false, FileDownloadError.FileNotAvailable);
			}
			lock (StateLock)
			{
				Reset(state);
			}
		}

		#endregion

		#region Hash Tree Download

		private bool DownloadHashTree(UserConnection connection, SourceInfo source)
		{
			log.Info("Downloading hash tree from " + source.User, null, this);
			if (connection.Supports(RequestType.HashTree))
			{
				Request request = new HashTreeRequest(remoteFile.TTH);
				WorkerState state = new WorkerState(connection, source, null, new MemoryStream());

				connections.Add(connection, state);
				try
				{
					connection.Download(request, HandleHashTreeData, state);
					connection.TransferAborted += OnTransferAborted;

					State = DownloadState.Active;
					return true;
				}
				catch (Exception e)
				{
					log.Error("Error downloading hash tree from " + source.User, e, this);
					Reset(state);
					throw;
				}
			}
			else
			{
				log.Error("Client to old: " + source.User, null, this);
				SetValid(source, false, FileDownloadError.ClientToOld);
			}
			return false;
		}

		private void HandleHashTreeData(byte[] buffer, int offset, int count,
			long remaining, object s)
		{
			WorkerState state = (WorkerState)s;
			try
			{
				state.Stream.Write(buffer, offset, count);
			}
			catch (Exception e)
			{
				StopWithException(e);
				throw;
			}
			if (remaining == 0)
			{
				try
				{
					hashTree = LoadHashTree(state.Stream, remoteFile);
					hashStore.AddHashTree(hashTree);

					log.Info("Hash tree download complete. Leaf count: " +
						hashTree.LeafCount, null, this);

					lock (StateLock)
					{
						CheckTempFile();
					}
				}
				catch (IOException e)
				{
					log.Error("Error loading the hash tree.", e, this);
					SetValid(state.Source, false, FileDownloadError.InvalidHashTree);
				}
				lock (StateLock)
				{
					Reset(state);
				}
			}
		}

		private static HashTree LoadHashTree(Stream stream, RemoteFileInfo remoteFile)
		{
			Tiger algorithm = Tiger.Create();
			long leafCount = stream.Length / algorithm.HashSize;

			long blockSize = HashTree.SegmentSize;
			while (blockSize * leafCount < remoteFile.Size)
			{
				blockSize *= 2;
			}

			stream.Position = 0;

			byte[] root = Base32.Decode(remoteFile.TTH);
			HashTree hashTree = new HashTree(Tiger.Create(), blockSize, remoteFile.Size, stream, root);
			hashTree.CalculateRoot();

			if (!HashTree.LeavesAreEqual(hashTree.Root, root))
			{
				throw new IOException("Invalid hash tree received.");
			}
			return hashTree;
		}

		#endregion

		#region File Download

		private bool DownloadFile(UserConnection connection, SourceInfo source)
		{
			Segment segment = segmentStrategy.CreateSegment(segmentInfo);
			if (segment == null)
			{
				log.Info("Not downloading from " + source.User + ". No segments available", null, this);
				return false;
			}
			UpdateBitArray(segmentInfo.Active, segment, true);

			log.Info("Downloading " + segment.Start + " to " + segment.End + " from " + source.User,
				null, this);

			source.Active = true;
			source.Monitor = connection.Monitor;

			WorkerState state = new WorkerState(connection, source, segment, null);

			// Open the output stream
			try
			{
				lock (streamLock)
				{
					if (outputStream == null)
					{
						OpenOutputStream(FileAccess.Write);
					}
				}
				state.Stream = new ValidatingOutputStream(Tiger.Create(),
					hashTree, outputStream, segment.Start);

				connections.Add(connection, state);
			}
			catch (Exception e)
			{
				log.Error("Error opening the output stream.", e, this);
				StopWithException(e);
				Reset(state);
				return false;
			}

			// Download the file
			try
			{
				Request request = new FileRequest(null, remoteFile.TTH,
					segment.Start, segment.End - segment.Start);

				connection.Download(request, HandleData, state);
				connection.TransferAborted += OnTransferAborted;

				State = DownloadState.Active;
				return true;
			}
			catch (Exception e)
			{
				log.Error("Error downloading from " + source.User, e, this);

				Reset(state);
				throw;
			}
		}

		private void HandleData(byte[] buffer, int offset, int count,
			long totalRemaining, object c)
		{
			WorkerState state = (WorkerState)c;
			Segment segment = state.Segment;
			try
			{
				if (segment.Position + count > segment.End)
				{
					count = (int)(segment.End - segment.Position);
				}

				ValidatingOutputStream stream = (ValidatingOutputStream)state.Stream;

				// Validate and write to file
				int validated = stream.ValidatedLeaves;
				lock (streamLock)
				{
					outputStream.Position = segment.Position;
					stream.Write(buffer, offset, count);

					// Update the bitfield
					for (; validated < stream.ValidatedLeaves; validated++)
					{
						segmentInfo.Completed[stream.StartLeaf + validated] = true;
					}

					UpdateDownloadPosition(count);
					UpdateMonitor(count);
				}

				segment.Position += count;
				if (segment.Position == segment.End)
				{
					log.Debug("Downloaded " + segment.Start + " - " + segment.End + " from " + state.Source.User,
						null, this);

					lock (StateLock)
					{
						Reset(state);
						if (downloadedBytes == remoteFile.Size)
						{
							HandleDownloadComplete();
						}
					}
				}
			}
			catch (Exception e)
			{
				log.Error("Error writing data from " + state.Source.User, e, this);
				if (e is ValidationException)
				{
					SetValid(state.Source, false, FileDownloadError.ValidationError);
				}
				else
				{
					StopWithException(e);
				}
				throw; // So that the connection gets to handle the exception and disconnect.
			}
		}

		// Hold before calling: StateLock
		private void HandleDownloadComplete()
		{
			log.Info("Download complete", null, this);

			string tempFile = outputStream.Name;
			string file = Path.Combine(savePath, remoteFile.Name);

			lock (streamLock)
			{
				CloseOutputStream();
			}
			try
			{
				if (!Directory.Exists(savePath))
				{
					Directory.CreateDirectory(savePath);
				}
				if (File.Exists(file))
				{
					File.Delete(file);
				}
				File.Move(tempFile, file);
			}
			catch (Exception e)
			{
				log.Error("Could not move temp file to the real location.");
				Exception = e;
			}
			State = DownloadState.Done;
		}

		#endregion

		#region FileStream

		// Hold before calling: streamLock
		private void OpenOutputStream(FileAccess fileAccess)
		{
			if (!Directory.Exists(savePath))
			{
				Directory.CreateDirectory(savePath);
			}
			outputStream = new FileStream(GetTempFileName(), FileMode.OpenOrCreate,
				fileAccess, FileShare.Read);
		}

		 // Hold before calling: streamLock
		private void CloseOutputStream()
		{
			if (outputStream != null)
			{
				outputStream.Close();
				outputStream = null;
			}
		}

		private string GetTempFileName()
		{
			return Path.Combine(savePath, remoteFile.Name + ".dctmp");
		}

		#endregion

		private static void SetValid(SourceInfo source, bool valid,
			FileDownloadError error)
		{
			source.IsValid = valid;
			source.Error = error;
		}

		private void UpdateBitArray(BitArray array, Segment segment, bool value)
		{
			int start = (int)(segment.Start / hashTree.BlockSize);
			int end = segment.End == hashTree.FileSize ? hashTree.LeafCount :
				(int)(segment.End / hashTree.BlockSize);

			for (int i = start; i < end; i++)
			{
				array[i] = value;
			}
		}	

		#endregion
	}
}
